# %%
from conn import clip_annotated_record, ac_task_col, check_frame_label_col, mongo_client
from tqdm import tqdm

# %%
anno_type = 21
clip_ids = []

with open("./clip_ids_21.txt", "r") as fp:
    for line in fp:
        clip_ids.append(line.strip())
print(len(clip_ids))

# %%
CHUNK = 1000

clip_id_chunks = [clip_ids[i : i + CHUNK] for i in range(0, len(clip_ids), CHUNK)]

print(len(clip_id_chunks))

# %%
res = ac_task_col.aggregate(
    [
        {"$match": {"annotate_type": anno_type}},
        {"$project": {"size": {"$size": "$clip_ids"}}},
        {"$group": {"_id": None, "total": {"$sum": "$size"}}},
    ]
)
for item in res:
    print(item)
# 累计 54999 帧


# %%
clip_id_set = set(clip_ids)


# 更新任务状态，删除验收 frame 记录，更新 clip 状态
def update_one_ac_task(task):
    cur_clip_ids = task["clip_ids"]
    cur_clip_info_map = {
        info["clip_id"]: info for info in task["clip_annotated_records"]
    }

    print("--- origin clip_ids_num", len(cur_clip_ids))

    left_clip_ids = set(cur_clip_ids) - clip_id_set
    if len(left_clip_ids) == 0:
        print("delete task ", task["_id"])
        ac_task_col.delete_one({"_id": task["_id"]})
    else:
        new_clip_ids_num = len(left_clip_ids)
        new_clip_ids = list(left_clip_ids)
        new_clip_annotated_records = [
            cur_clip_info_map[clip_id] for clip_id in new_clip_ids
        ]
        print("update task ", task["_id"], " new_clip_ids_num: ", new_clip_ids_num)
        ac_task_col.update_one(
            {"_id": task["_id"]},
            {
                "$set": {
                    "clip_ids": new_clip_ids,
                    "clip_ids_num": new_clip_ids_num,
                    "clip_annotated_records": new_clip_annotated_records,
                }
            },
        )


# %%

cnt = 0
cur_tasks = ac_task_col.find(
    {"clip_ids": {"$in": clip_ids}, "annotate_type": anno_type},
    {"status": 1, "clip_ids": 1, "clip_ids_num": 1, "clip_annotated_records": 1},
).to_list()
for cur_task in cur_tasks:
    update_one_ac_task(cur_task)


# %%

with mongo_client.start_session() as session, session.start_transaction():
    for clip_id_chunk in clip_id_chunks:
        clip_annotated_record.update_many(
            {"clip_id": {"$in": clip_id_chunk}, "annotate_type": anno_type},
            {"$set": {"accept_status": None}},
            session=session,
        )
